package main.java.lpr;

import java.util.ArrayList;
import java.util.List;

import ai.djl.modality.cv.translator.BaseImageTranslator;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.types.DataType;
import ai.djl.translate.TranslatorContext;

/**
 * LprE2eTranslator
 * @date: 2022-02-25 14:55
 * @author: alvin
 * @copyright: Copyright © 2022 YZZH. All rights reserved.
 * @description: 端到端车牌识别
 */
public class LprE2eTranslator  extends BaseImageTranslator<PlateChar>{

    private List<String> classes;


    public List<String> getClasses() {
        return classes;
    }

    public void setClasses(List<String> classes) {
        this.classes = classes;
    }

    /**
     * Constructs an Image Classification using {@link Builder}.
     *
     * @param builder the data to build with
     */
    public LprE2eTranslator(Builder builder) {
        super(builder);
        this.classes = builder.getOptClasses();
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public PlateChar processOutput(TranslatorContext ctx, NDList list) {
        NDArray array = list.singletonOrThrow();
        NDArray array2 = array.toType(DataType.FLOAT64, false);
        NDArray res = array2.argMax(1);
        res = res.toType(DataType.INT32, false);
        int[] labs = res.toIntArray();
        String results = "";
        Float confidence = (float)0.0;
        List<Float>  charConfidences = new ArrayList<>();
        for (int i = 0; i < labs.length; i++) {
            int one = labs[i];
            if (one < classes.size() && (i == 0 || (one != labs[i - 1]))) {
                results += classes.get(one);
                Float confi = array.get(i).get(one).toArray()[0].floatValue();
                charConfidences.add(confi);
                confidence += confi;
            } 
        }
        if (results.length() != 0) {
            confidence /= results.length();
        } 
        PlateChar plateChar = new PlateChar();
        plateChar.setPlateNumber(results);
        plateChar.setConfidence(confidence);
        plateChar.setCharConfidences(charConfidences);

        return plateChar;
    }

    

    /**
     * Creates a builder to build a {@code ImageClassificationTranslator}.
     *
     * @return a new builder
     */
    public static Builder builder() {
        return new Builder();
    }


    /**
     * A Builder to construct a {@code ImageClassificationTranslator}.
     */
    public static class Builder extends ClassificationBuilder<Builder> {

        
        private List<String> optClasses; 


        public List<String> getOptClasses() {
            return optClasses;
        }

        public void setOptClasses(List<String> optClasses) {
            this.optClasses = optClasses;
        }

        Builder() {

        }

        /**
         * {@inheritDoc}
         */
        @Override
        protected Builder self() {
            return this;
        }


        public LprE2eTranslator build() {
            validate();
            return new LprE2eTranslator(this);
        }



        public Builder optClasses(List<String> classes) {
            this.optClasses = classes;
            return this;
        }
    }    
}


